package tree

//红黑树泛型实现
//参考资料：https://www.cnblogs.com/gcheeze/p/11186806.html
//特性：
//(1) 每个节点或者是黑色，或者是红色。
//(2) 根节点是黑色。
//(3) 每个叶子节点是黑色。 [注意：这里叶子节点，是指为空的叶子节点！]
//(4) 如果一个节点是红色的，则它的子节点必须是黑色的。
//(5) 从一个节点到该节点的子孙节点的所有路径上包含相同数目的黑节点。

type Rbtree[K, T any] interface {
	Key(v T) K        //节点的key
	Less(l, r K) bool //节点排序规则
	Len() int
	Replace(v T) *RbtreeNode[T]          //插入节点，重复则替换
	Insert(v T) *RbtreeNode[T]           //插入节点，重复则忽视
	InsertRepeatable(v T) *RbtreeNode[T] //插入节点，允许数据重复
	Front() *RbtreeNode[T]               //查询有序节点的第一个节点
	Back() *RbtreeNode[T]                //查询有序节点的最后一个节点
	Find(key K) *RbtreeNode[T]           //查找节点，如果存在重复，则返回第一个找到的节点
	LowerBound(key K) *RbtreeNode[T]     //查找第一个不小于(>=)的节点
	UpperBound(key K) *RbtreeNode[T]     //查找第一个大于的节点
	Erase(n *RbtreeNode[T])              //删除节点
	Remove(key K) (T, bool)              //删除一个相同的节点
	RemoveAll(key K) []T                 //删除所有相同的节点
}

//定义一个红黑树
//getKey:获取value用于比较的key
//lessthan:定义key的排序规则
func NewRbtree[K, T any](getKey func(T) K, lessthan func(K, K) bool) Rbtree[K, T] {
	tree := &rbtree[K, T]{
		len:  0,
		key:  getKey,
		less: lessthan,
	}
	sentinel := &RbtreeNode[T]{
		red: false,
		//Value: "sentinel",
	}
	sentinel.t = sentinel
	sentinel.p = sentinel
	sentinel.l = sentinel
	sentinel.r = sentinel
	tree.root = sentinel
	return tree
}

//定义红黑树结构
type rbtree[K, T any] struct {
	root *RbtreeNode[T]  //sentinel:parent对应根节点,left与right对应自身
	len  int             //当前存储的节点个数
	key  func(T) K       //获取value的key
	less func(K, K) bool //排序规则
}

//定义红黑树节点结构
type RbtreeNode[T any] struct {
	t     *RbtreeNode[T] //sentinel
	p     *RbtreeNode[T] //parent
	l     *RbtreeNode[T] //左节点
	r     *RbtreeNode[T] //右节点
	red   bool           //是红色还是黑色
	Value T              //值
}

//查找左子树的最小值
func (n *RbtreeNode[T]) minimum() *RbtreeNode[T] {
	//fmt.Printf("%v %v %v %v %v %v\n", n.t.p.Value, n.p.Value, n.l.Value, n.r.Value, n.red, n.Value)
	s := n
	l := s.l
	for l != n.t {
		s = l
		l = l.l
	}
	return s
}

//查找右子树的最大值
func (n *RbtreeNode[T]) maximum() *RbtreeNode[T] {
	s := n
	r := s.r
	for r != n.t {
		s = r
		r = r.r
	}
	return s
}

//查询前序节点
func (n *RbtreeNode[T]) predecessor() *RbtreeNode[T] {
	if n.l != n.t {
		//n拥有左子树,前序节点一定是左子树的最大节点
		return n.l.maximum()
	} else {
		// 找到高层节点中以p所在的树为右子树的树的根节点，即是前驱节点
		node := n
		p := node.p
		for p != n.t && p.l == node {
			//node是左节点，其父节点比node大
			node = p
			p = node.p
		}
		if p != n.t {
			return p //p.r == node
		}
	}
	return nil
}

//查询后继节点
func (n *RbtreeNode[T]) successor() *RbtreeNode[T] {
	if n.r != n.t {
		//n拥有右子树，后继节点一定是右子树中的最小节点
		return n.r.minimum()
	} else {
		// 找到高层节点中以p所在的树为左子树的树的根节点，即是后继节点
		node := n
		p := node.p
		for p != n.t && p.r == node {
			//node是右节点，其父节点比node小
			node = p
			p = node.p
		}
		if p != n.t {
			return p //p.l == node
		}
		//else 没有比n更大的节点
	}
	return nil
}

//前一个最接近的节点
func (n *RbtreeNode[T]) Prev() *RbtreeNode[T] { return n.predecessor() }

//下一个最接近的节点
func (n *RbtreeNode[T]) Next() *RbtreeNode[T] { return n.successor() }

func (t *rbtree[K, T]) Key(v T) K {
	return t.key(v)
}

func (t *rbtree[K, T]) Less(l, r K) bool {
	return t.less(l, r)
}

//当前节点个数
func (t *rbtree[K, T]) Len() int {
	return t.len
}

//判断节点的值是否相等
func (t *rbtree[K, T]) equal(l, r K) bool {
	return !t.Less(l, r) && !t.Less(r, l)
}

//顶点
func (t *rbtree[K, T]) top() *RbtreeNode[T] { return t.root.p }

//有序节点的第一个节点
func (t *rbtree[K, T]) Front() *RbtreeNode[T] {
	if t.len == 0 {
		return nil
	}
	return t.root.p.minimum()
}

//有序节点的最后一个节点
func (t *rbtree[K, T]) Back() *RbtreeNode[T] {
	if t.len == 0 {
		return nil
	}
	return t.root.p.maximum()
}

//插入节点，重复则替换
func (t *rbtree[K, T]) Replace(v T) *RbtreeNode[T] {
	return t.insert(v, false, true)
}

//插入节点，重复则忽视
func (t *rbtree[K, T]) Insert(v T) *RbtreeNode[T] {
	return t.insert(v, false, false)
}

//插入节点，允许重复
func (t *rbtree[K, T]) InsertRepeatable(v T) *RbtreeNode[T] {
	return t.insert(v, true, false)
}

func (t *rbtree[K, T]) Find(key K) *RbtreeNode[T] {
	p := t.top()
	for p != t.root {
		if t.Less(key, t.Key(p.Value)) {
			p = p.l
		} else if t.Less(t.Key(p.Value), key) {
			p = p.r
		} else {
			return p
		}
	}
	return nil
}

//查找第一个不小于(>=)的节点
func (t *rbtree[K, T]) LowerBound(key K) *RbtreeNode[T] {
	p := t.top()
	var where *RbtreeNode[T]
	for p != t.root {
		if t.Less(t.Key(p.Value), key) {
			//比v小，则找右边
			p = p.r
		} else {
			//不比v小，则找左边是否存在跟解决的项，或者找到第一个相等的项
			where = p
			p = p.l
		}
	}
	return where
}

//查找第一个大于的节点
func (t *rbtree[K, T]) UpperBound(key K) *RbtreeNode[T] {
	p := t.top()
	var where *RbtreeNode[T]
	for p != t.root {
		if t.Less(key, t.Key(p.Value)) {
			//比v小，则找右边
			where = p
			p = p.l //尝试比当前更小的项
		} else {
			p = p.r
		}
	}
	return where
}

//删除项
func (t *rbtree[K, T]) Erase(n *RbtreeNode[T]) {
	if n.t == t.root {
		t.remove(n)
	}
}

//删除第一个相同的节点
func (t *rbtree[K, T]) Remove(key K) (T, bool) {
	p := t.Find(key)
	if p != nil {
		t.remove(p)
		return p.Value, true
	}
	var r T
	return r, false
}

//删除所有相同的节点
func (t *rbtree[K, T]) RemoveAll(key K) []T {
	var r []T
	for {
		if val, ok := t.Remove(key); ok {
			r = append(r, val)
		} else {
			break
		}
	}
	return r
}

//插入节点
func (t *rbtree[K, T]) insert(v T, repeatable bool, replace bool) *RbtreeNode[T] {
	key := t.Key(v)
	//查找插入位置
	p := t.top()
	if p == t.root {
		//插入根节点
		node := &RbtreeNode[T]{
			t:     t.root,
			p:     t.root,
			l:     t.root,
			r:     t.root,
			red:   false, //黑色
			Value: v,
		}
		t.root.p = node
		t.len = 1
		return node
	}
	//查找插入位置
	findp := p
	addleft := false
	if repeatable {
		//允许重复插入
		for findp != t.root {
			p = findp
			if t.Less(key, t.Key(p.Value)) {
				//v<p,继续找左边
				findp = p.l
				addleft = true
			} else {
				//v>=p,继续找右边,重复项插入最右边
				findp = p.r
				addleft = false
			}
		}
	} else {
		//不允许重复插入
		for findp != t.root {
			p = findp
			if t.Less(key, t.Key(p.Value)) {
				//v<p,继续找左边
				findp = p.l
				addleft = true
			} else {
				//v>=p
				if !t.Less(t.Key(p.Value), key) {
					//v<>p,替换值
					if replace {
						p.Value = v
					}
					return p
				}
				//v>p
				findp = p.r
				addleft = false
			}
		}
	}
	node := &RbtreeNode[T]{
		t:     t.root,
		p:     p,
		l:     t.root,
		r:     t.root,
		red:   true,
		Value: v,
	}
	if addleft {
		p.l = node
	} else {
		p.r = node
	}
	t.len++
	t.fixup_insert(node)
	return node
}

//左旋:z变为右节点的左子节点。右节点变为父节点，z节点变为右节点的左子节点
func (t *rbtree[K, T]) rotate_left(z *RbtreeNode[T]) {
	rp := z.r //右节点变为父节点
	rp.p = z.p
	//改变父节点
	if z.p == t.root {
		t.root.p = rp
	} else {
		if z == z.p.l {
			//如果z是它父节点的左孩子，则将右节点设为该父节点的左孩子
			z.p.l = rp
		} else {
			z.p.r = rp
		}
	}
	//右节点的左子树变为该节点的右子树
	z.r = rp.l
	if rp.l != t.root {
		rp.l.p = z
	}
	rp.l = z
	z.p = rp
}

//右旋:z变为左节点的右子节点。左节点变为父节点，z节点变为左节点的右子节点
func (t *rbtree[K, T]) rotate_right(z *RbtreeNode[T]) {
	rp := z.l //左节点变为父节点
	rp.p = z.p
	//改变父节点
	if z.p == t.root {
		t.root.p = rp
	} else {
		if z == z.p.l {
			//如果z是它父节点的左孩子，则将左节点设为该父节点的左孩子
			z.p.l = rp
		} else {
			z.p.r = rp
		}
	}
	//左节点的右子树变为该节点的左子树
	z.l = rp.r
	if rp.r != t.root {
		rp.r.p = z
	}
	rp.r = z
	z.p = rp
}

//插入的节点是红色，只有当其父节点是红色时违背红黑树规则
func (t *rbtree[K, T]) fixup_insert(z *RbtreeNode[T]) {
	//z.red == true
	for z.p != t.root && z != t.root && z.p.red {
		//父节点为红色，必存在z.p.p
		if z.p == z.p.p.l {
			//父节点是左子树
			s := z.p.p.r //z节点的叔节点
			if s != t.root && s.red {
				//叔节点是红色
				//祖父节点的子节点变黑，祖父节点变红
				z.p.red = false
				s.red = false
				z.p.p.red = true
				//继续调整祖父
				z = z.p.p
			} else {
				//叔节点是黑色,或者是sentinel
				if z == z.p.r {
					//z >= z.p
					//通过左旋把z变成父节点，z.p变成z的左子节点
					z = z.p
					t.rotate_left(z)
					//'.左旋变成.',z还是在最低层的.,从右边变成左边
				}
				z.p.red = false
				z.p.p.red = true
				t.rotate_right(z.p.p)
				//.-'右旋-'-,z变成左边的-,其中'为z.p黑色，右边的-为z.p.p红色，未改变黑高度
			}
		} else {
			//父节点是右子树
			s := z.p.p.l //叔节点
			if s != t.root && s.red {
				//叔节点是红色，与左子树模式相同
				z.p.red = false
				s.red = false
				z.p.p.red = true
				z = z.p.p
			} else {
				//叔节点是黑色,或者是sentinel
				if z == z.p.l {
					//通过右旋把z变成父节点，z.p变成z的左子节点
					z = z.p
					t.rotate_right(z)
					//.'右旋变成'.,z还是在最底层的.,从左边变成右边
				}
				z.p.red = false
				z.p.p.red = true
				t.rotate_left(z.p.p)
				//'-.右旋-'-,z变成右边的-,其中'为z.p黑色，左边的-为z.p.p红色，未改变黑高度
			}
		}
	}
	t.root.p.red = false //设置顶点为黑色
}

//删除一个节点，总共分三种情况
// 1.被删除节点没有儿子，即为叶节点。那么，直接将该节点删除就OK了。
// 2.被删除节点只有一个儿子。那么，直接删除该节点，并用该节点的唯一子节点顶替它的位置。
// 3.被删除节点有两个儿子。那么，先找出它的后继节点；然后把“它的后继节点的内容”复制给“该节点的内容”；
//之后，删除“它的后继节点”。在这里，后继节点相当于替身，在将后继节点的内容复制给"被删除节点"之后，再将后继节点删除。
//这样就巧妙的将问题转换为"删除后继节点"的情况了，下面就考虑后继节点。 在"被删除节点"有两个非空子节点的情况下，它的后继节点不可能是双子非空。
//既然"它的后继节点"不可能双子都非空，就意味着"该节点的后继节点"要么没有子节点，要么只有一个子节点。
//若没有子节点，则按"情况1"进行处理；若只有一个子节点，则按"情况2"进行处理。
func (t *rbtree[K, T]) remove(n *RbtreeNode[T]) {
	d := n      //顶替删除的节点,即实际删除的节点
	x := t.root //调整节点，待删除的位置
	//node
	if n.l == t.root || n.r == t.root {
		//情况1或2
		d = n
	} else {
		d = n.successor() //情况3
	}
	if d.l != t.root {
		x = d.l      //树中删除除位置
		d.l = t.root //删除
	} else if d.r != t.root {
		x = d.r
		d.r = t.root
	}
	t.len--
	if x != t.root {
		//x转移到d，相当于删掉x位置
		x.p = d.p //子节点顶替父节点
		if d.p == t.root {
			t.root.p = x //删除节点是跟节点，则重置根节点
		} else if d == d.p.l {
			d.p.l = x
		} else {
			d.p.r = x
		}
		//x.l = x.r = t.root
		x.red = d.red
		//单叶子节点必为红色节点
		x = t.root //删除的是红色节点，全树平衡
	} else {
		//没有x，直接删除d位置
		if d == t.root.p {
			//顶点
			t.root.p = t.root
			t.len = 0
			return
		}
		x = d //调整
	}
	//调整x节点
	if x != t.root {
		//经过x的子树黑色节点少1
		//因此兄弟节点的子树节点也要减1
		t.fixup_del(x)
		//删除x
		if x == x.p.l {
			x.p.l = t.root
		} else if x == x.p.r {
			x.p.r = t.root
		}
	}
	if d != n {
		//使用d顶替node，复制d到node
		n.Value = d.Value
		//t.moveNode(d, n)
	}
}

func (t *rbtree[K, T]) fixup_del(z *RbtreeNode[T]) {
	//如果当前节点是红色直接变成黑色即可保持特性5成立
	for z != t.root && z != t.root.p && !z.red {
		//经过x的路径少一个黑色
		if z == z.p.l {
			//当前为左节点
			w := z.p.r //兄弟节点
			if w.red { //红色兄弟节点
				w.red = false      //兄弟节点变黑
				z.p.red = true     //父节点变红
				t.rotate_left(z.p) //w变成祖父节点
				w = z.p.r          //当前改变后的兄弟节点
			}
			if !w.l.red && !w.r.red {
				w.red = true //兄弟节点变红
				z = z.p
			} else {
				if !w.r.red {
					w.l.red = false
					w.red = true
					t.rotate_right(w)
					w = z.p.r //当前改变后的兄弟节点
				}
				w.red = z.p.red
				z.p.red = false
				w.r.red = false
				t.rotate_left(z.p)
				z = t.top()
			}
		} else {
			w := z.p.l
			if w.red {
				w.red = false
				z.p.red = true
				t.rotate_right(z.p)
				w = z.p.l
			}
			if !w.l.red && !w.r.red {
				w.red = true
				z = z.p
			} else {
				if !w.l.red {
					w.r.red = false
					w.red = true
					t.rotate_left(w)
					w = z.p.l
				}
				w.red = z.p.red
				z.p.red = false
				t.rotate_right(z.p)
				z = t.top()
			}
		}
	}
	z.red = false
}

func (t *rbtree[K, T]) moveNode(d, n *RbtreeNode[T]) {
	d.p = n.p
	if n.p == t.root {
		t.root.p = d
	} else if n == n.p.l {
		n.p.l = d
	} else if n == n.p.r {
		n.p.r = d
	}
	if n.l != d {
		d.l = n.l
		if n.l != t.root {
			n.l.p = d
		}
	} else {
		d.l = t.root
	}
	if n.r != d {
		d.r = n.r
		if n.r != t.root {
			n.r.p = d
		}
	} else {
		d.r = t.root
	}
	d.red = n.red
}
